import json
import time
import base64
import hashlib
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5, AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
import Crypto
import requests


class Aes:
    def __init__(self, key: str, iv: str = 'a46ef9a97bb04ea6'):
        self.key = key.encode()
        self.iv = iv.encode()

    def encrypt(self, text):
        cryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
        pad_pkcs7 = pad(text.encode('utf-8'), AES.block_size, style='pkcs7')
        ciphertext = cryptor.encrypt(pad_pkcs7)
        ciphertext = ciphertext.hex()
        return ciphertext

    def decrypt(self, text):
        cryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
        text = bytes.fromhex(text)
        plain_text = cryptor.decrypt(text)
        plain_text = unpad(plain_text, AES.block_size, style='pkcs7')
        return plain_text.decode()


class Rsa:
    MAX_ENCRYPT_BLOCK = 117
    MAX_DECRYPT_BLOCK = 128

    @classmethod
    def encrypt(cls, text, public_key):
        key = RSA.import_key(public_key)
        plaintext = text.encode()
        ciphertext = b''
        for i in range(0, len(plaintext), cls.MAX_ENCRYPT_BLOCK):
            chunk = plaintext[i:i + cls.MAX_ENCRYPT_BLOCK]
            cipher = PKCS1_v1_5.new(key)
            encrypted_chunk = cipher.encrypt(chunk)
            ciphertext += encrypted_chunk
        ciphertext = base64.b64encode(ciphertext).decode()
        return ciphertext

    @staticmethod
    def decrypt(text, private_key):
        pass

    @classmethod
    def verify_signature(cls, signature, data, public_key):
        from Crypto.Signature import pkcs1_15
        from Crypto.Hash import SHA256
        key = RSA.import_key(public_key)
        msg_hash = SHA256.new(data)
        cipher = pkcs1_15.new(key)
        cipher.verify(msg_hash, signature)


class UVerif:
    def __init__(self, root, apikey, appid, version='default', public_key=None):
        self.apikey = apikey
        self.api = f'{root}/api/{appid}/{version}'
        self.public_key = public_key
        self.token = None

    def handle_res(self, res, aes_key):
        try:
            res_json = res.json()
        except:
            res_json = {}
        if res_json.get('code') != 200:
            raise Exception(res.text)
        if self.public_key:
            # RSA验签
            sign = res_json['sign'].encode()
            sign = base64.b64decode(sign)
            data = res_json['data'].encode()
            Rsa.verify_signature(sign, data, self.public_key)
            # AES解密
            res_data = self.decrypt(res_json, aes_key.hex())
        else:
            res_data = json.loads(res_json['data'])
        return res_data

    def login(self, kami, password=None):
        aes_key = get_random_bytes(16)
        data = {
            "user": kami,
            "password": password,
            "udid": "005FE82DFC9EA5B4A37",
            "time": int(time.time()),
            "key": aes_key.hex()
        }
        url = f'{self.api}/logon'
        res = self.request(url, data=data, aes_key=aes_key)
        if res.get('token'):
            self.token = res.get('token')
        return res

    @staticmethod
    def body(data):
        sorted_data = dict(sorted(data.items(), key=lambda x: x[0]))
        filtered_data = {k: v for k, v in sorted_data.items() if v is not None}
        url_encoded_data = '&'.join([f"{k}={v}" for k, v in filtered_data.items()])
        return url_encoded_data

    def request(self, url, data, aes_key):
        body = self.body(data)
        sign = self.md5_sign(body)
        if self.public_key:
            body = Rsa.encrypt(body, self.public_key)
            data = {
                "data": body
            }
        data['sign'] = sign
        res = requests.post(url=url, data=data)
        res = self.handle_res(res, aes_key)
        return res

    def verif(self):
        aes_key = get_random_bytes(16)
        data = {
            "token": self.token,
            "time": int(time.time()),
            "key": aes_key.hex()
        }
        url = f'{self.api}/verif'
        res = self.request(url, data=data, aes_key=aes_key)
        self.token = res['token']
        print('验证token', res)

    def md5_sign(self, data):
        return hashlib.md5((data + self.apikey).encode()).hexdigest()

    def encrypt(self, key):
        pass

    def decrypt(self, res_json, key):
        if self.public_key:
            aes = Aes(key)
            res = aes.decrypt(res_json['data'])
            return json.loads(res)
        return res_json


if __name__ == '__main__':
    # 1、客户端生成一个AES对称密钥，放到请求体中
    # 2、使用RSA公钥加密所有请求体，客户端发送请求
    # 3、服务器解密请求体，进行业务处理
    # 4、服务器使用本次请求的AES密钥加密响应体
    # 5、服务器使用应用的RSA私钥签名AES加密后的数据，返回响应
    # 6、客户端收到响应，使用RSA公钥验证加密数据
    # 7、验证通过后使用本次请求的AES密钥解密数据获得业务数据进行业务处理

    # 签名主要是验证数据是服务器返回的，而不是通过中间人篡改，HTTP响应注入的。
    # AES加密主要是保证数据传输过程中不会被看到数据结构，修改响应数据。
    public_key = """-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDNM0WAPF74oUlGSpys44oU9aUH
1OxrNDVHdtIzg9P1mO7V0x8AtVCZrM+wzHR9Z/jDkTgKa9A9L9OeMzxBC+vTr1ij
LMklQKLvUoHG/AlZpKtmflYQsfhysTf8AaFQFsDuy4r+m9WcX/P9acf/GtOmFVeG
hSqJTHI9h7JEXicgrQIDAQAB
-----END PUBLIC KEY-----"""
    client = UVerif(root='http://ukami.com', apikey='77415bf8b3d51a24f6fe6a65efa179aa', appid=2,
                    public_key=public_key)
    client.login('8605BEB67EF63BCWIQ')
    for i in range(10):
        client.verif()
        time.sleep(10)
